import meep as mp                                  ###载入meep库
import numpy as np                                 ###载入numpy库                                                   
import matplotlib.pyplot as plt                    ###载入python绘图工具
from meep.materials import Au                      ###从MEEP的材料库中导入所需的材料
Ix = []                                            ###主像元能量列表初始化
Ix_adj = []                                        ###邻近像元能量列表初始化
wavelength = 10.55                                 ###入射波长
period = np.linspace(5, 10, 5)                     ###量子阱周期
print(period)                                      ###查看量子阱周期是否正确
###定义折射率生成函数
### START CODE HERE
def refractive_index(wavelength, x_Al, fitted_constant):      
    """
    This funtion is used to compute the refractive index for arbitray incident wavelength and fraction of Al
    Arugements:
       Input:
             wavelength -- incident wavelength of the Source
             x_Al -- fraction of the Al in AlxGa1-xAs
             fitted_constant -- the bias constant to correct the refractive index
       Output:
             n -- real part of the refractive index
    """
    q = 1.6e-19                                        ### unit Charge
    E0 = (1.425 + 1.155 * x_Al + 0.37*x_Al**2)*q       ### fundamental band gap at gamma-point
    E0_plus_delta = (1.765 + 1.115 * x_Al + 0.37*x_Al**2)*q     ### spin-orbit splitting energy
    A0 = 6.3 + 19.0*x_Al - fitted_constant
    B0 = 9.4 - 10.2*x_Al - fitted_constant
    c = 2.998e8                                            ### light speed
    h = 6.625e-34                                          ### Plank Constant
    kai_s0 = h*c/wavelength/(E0_plus_delta)
    kai = h*c/wavelength/E0
    f_kai = f(kai)
    f_kai0 = f(kai_s0)
    n = (A0*(f_kai+f_kai0/2*(E0/E0_plus_delta)**1.5)+B0)**0.5
    
    return n
def f(kai):
    """
    This function is used to compute the middle varible for the computation of the refractive index
    """
    f = (2-(1-kai)**0.5-(1+kai)**0.5)/kai**2
    return f
### END CODE HERE
### 计算相对介电常数
### START CODE HERE
def epsilon_index(x_Al):
    """
    Function to compute the dielectric constant for the arbitray fraction of Al.
    
    Arugement:
      Input:
        x_Al -- fraction of Al
      output:
        eps -- dielectric constant epsilon of the arbitray fraction of Al
    """
    eps = 12.90-2.84*x_Al
    return eps
### END CODE HERE
### 角度因子函数
### START CODE HERE
def tan_factor(angle):
    angle = angle*np.pi/180
    angle_factor = np.tan(angle)
    return angle_factor
### END CODE HERE
### 计算物理参数
n0 = refractive_index(wavelength=wavelength*1e-6, x_Al=0, fitted_constant=0.01)   ### GaAs的折射率
eps0 = epsilon_index(x_Al=0)                                                      ### GaAs的相对介电常数
n_co = refractive_index(wavelength = wavelength*1e-6, x_Al = 0.55, fitted_constant = 0.01)  ### Al0.55Ga0.45As的折射率
eps_co = epsilon_index(x_Al = 0.55)                                                         ### Al0.55Ga0.45As的相对介电常数
n_b = refractive_index(wavelength = wavelength*1e-6, x_Al = 0.18, fitted_constant = 0.01)   ### Al0.18Ga0.82As的折射率
eps_b = epsilon_index(x_Al = 0.18)                                                  ### Al0.18Ga0.82As的折射率
### 检查物理参数
print(n0, eps0)
print(n_co, eps_co)
print(n_b, eps_b)
### Main Loop program
Loop_num = 0
for periods in period:
    angle = 45                      ### reflective angle
    angle_factor = tan_factor(angle)### tan
    L = 15                          ### size of pixels
    structure = 'reflective'    ### 2 choices 'block' or 'reflective'
    ncells = 3                       ### number of the cells
    geometry = []                   # initialize the structure
    dpml = 1.0                     # thickness of the boundary layers
    t_refl = 1/15*L                # thickness of the metal        
    t_electrode = 1.0  # thickness of the electrode       
    l_electrode = 15   # length of the electrode
    t_well = 6e-3 # thickness of the Quantum well
    t_br = 45e-3  # thickness of the Barrier
    t_sub = 3.0   # thickness of the substrate
    periods = int(periods)  # periods of the Quantum Well
    t_top = 1.3/15*L   # thickness of the top contact from 1.3 to 0.4
    t_bot = 1.2/15*L # thickness of the bottom contact from 1.2 to 0.4
    t_co = 0.3/15*L    # thickness of the corrosion barrier 
    
    fwidth = 1    # spectrum width of the source
    source_width = (L - 2*t_refl)/15*L   # width of the source
    dair = (source_width)*(13)/2.44/10.55   # distance between the source and substrate
    print(dair)                          
    t_structure = periods*(t_well+t_br)+t_bot+t_top       ### thickness of the hole structure 
    resolution = 2/t_well                                 ### grids per 1um
    plane = 'xy'                                          ### 2 choices xy and yz, to determine the plane
    ### Cell size
    sx = 2*dpml + dair + t_sub + t_co + t_bot + periods*(t_well+t_br) + t_top + t_electrode  
    sy = 2*dpml + ncells*L
    sz = 2*dpml + ncells*L
    ### Generating the structure
    geometry.append(mp.Block(size = mp.Vector3(t_sub, ncells*L, L),
                             center = mp.Vector3(-sx/2+dpml+dair+t_sub/2),
                             material = mp.Medium(epsilon = eps0, index = n0)))  ## Material: GaAs
    geometry.append(mp.Block(size = mp.Vector3(t_co, ncells*L, L),
                             center = mp.Vector3(-sx/2+dpml+dair+t_sub+t_co/2),
                             material = mp.Medium(epsilon = eps_co, index = n_co)))   ## Material: Al0.55Ga0.45As
    
    geometry.append(mp.Block(size = mp.Vector3(t_bot, ncells*L, L-2*t_refl),
                             center = mp.Vector3(-sx/2+dpml+dair+t_sub+t_co+t_bot/2),
                             material = mp.Medium(epsilon = eps0, index = n0)))  ## Material: GaAs
    for i in range(periods):
        geometry.append(mp.Block(size = mp.Vector3(t_br, ncells*L, L-2*t_refl),
                                        center = mp.Vector3(sx/2-dpml-t_electrode-t_top-t_br/2-i*(t_br+t_well)),
                                        material = mp.Medium(epsilon = eps_b, index = n_b)))  ## Material: Al0.18Ga0.82A
        geometry.append(mp.Block(size = mp.Vector3(t_well, ncells*L, L-2*t_refl),
                                        center = mp.Vector3(sx/2-dpml-t_electrode-t_top-t_br-t_well/2-i*(t_br-t_well)),
                                        material = mp.Medium(epsilon = eps0, index = n0)))
    geometry.append(mp.Block(size = mp.Vector3(t_top, ncells*L, L-2*t_refl),
                             center = mp.Vector3(sx/2-dpml-t_electrode-t_top/2),  
                             material = mp.Medium(epsilon = eps0, index = n0)))   ## Material: GaAs
    if structure == 'block':
        ### Generate the metal
        for k in range(-ncells//2+1, ncells//2+1):
            geometry.append(mp.Block(size = mp.Vector3(t_structure, t_refl, t_structure), 
                             center = mp.Vector3(sx/2-dpml-t_electrode-t_structure/2, -sy/2+dpml+t_refl/2+k*L),
                             material = Au))
            geometry.append(mp.Block(size = mp.Vector3(t_structure, t_refl, t_structure),  
                             center = mp.Vector3(sx/2-dpml-t_electrode-t_structure/2, sy/2-dpml-t_refl/2+k*L),
                             material = Au))
            geometry.append(mp.Block(size = mp.Vector3(t_structure, L, t_refl),
                             center = mp.Vector3(x=sx/2-dpml-t_electrode-t_structure/2, y=0+k*L, z=L/2+t_refl/2),
                             material = Au))
            geometry.append(mp.Block(size = mp.Vector3(t_structure, L, t_refl),
                             center = mp.Vector3(x=sx/2-dpml-t_electrode-t_structure/2, y=0+k*L, z=-L/2-t_refl/2),
                             material = Au))
             ### Generate the electrode
            geometry.append(mp.Block(size = mp.Vector3(t_electrode, L, L),
                             center = mp.Vector3(sx/2-dpml-t_electrode/2, k*L),
                             material = Au))
            ### Generate the Grating
            num_gt = 5
            gt_p = (L - t_refl*2)/5
            gt_duty = 0.8
            gt_width = gt_p * gt_duty
            gt_length = t_structure / 6
            for i in range(-num_gt//2+1, num_gt//2+1):
                for j in range(-num_gt//2+1, num_gt//2+1):
                    geometry.append(mp.Block(size = mp.Vector3(gt_length, gt_width, gt_width),
                                             center = mp.Vector3(x=sx/2-dpml-t_electrode-gt_length/2,
                                                         y=i*gt_p+k*L,
                                                         z=j*gt_p),
                                             material = Au))
    elif structure == 'reflective':
        for j in range(-ncells//2+2, ncells//2+2):
            vertices1 = [mp.Vector3(x=-sx/2+dpml+dair+t_sub+t_co, y=sy/2-dpml-t_refl+j*L),
                     mp.Vector3(x=-sx/2+dpml+dair+t_sub+t_co+t_structure/2, y=sy/2-dpml-t_refl-t_structure/2/angle_factor+j*L),
                     mp.Vector3(x=sx/2-dpml-t_electrode, y=sy/2-dpml-t_refl-t_structure/angle_factor+j*L),
                     mp.Vector3(x=sx/2-dpml-t_electrode, y=sy/2-dpml-t_structure/angle_factor+j*L),
                     mp.Vector3(x=sx/2-dpml-t_electrode-t_structure/2, y=sy/2-dpml-t_structure/2/angle_factor+j*L),
                     mp.Vector3(x=sx/2-dpml-t_electrode-t_structure, y=sy/2-dpml+j*L)]
            vertices2 = [mp.Vector3(x=-sx/2+dpml+dair+t_sub+t_co, y=-(sy/2-dpml-t_refl)+j*L),
                     mp.Vector3(x=-sx/2+dpml+dair+t_sub+t_co+t_structure/2, y=-(sy/2-dpml-t_refl-t_structure/2/angle_factor)+j*L),
                     mp.Vector3(x=sx/2-dpml-t_electrode, y=-(sy/2-dpml-t_refl-t_structure/angle_factor)+j*L),
                     mp.Vector3(x=sx/2-dpml-t_electrode, y=-(sy/2-dpml-t_structure/angle_factor)+j*L),
                     mp.Vector3(x=sx/2-dpml-t_electrode-t_structure/2, y=-(sy/2-dpml-t_structure/angle_factor/2)+j*L),
                     mp.Vector3(x=sx/2-dpml-t_electrode-t_structure, y=-(sy/2-dpml)+j*L)]
            geometry.append(mp.Prism(vertices = vertices1,
                                 height = L,
                                 material = Au))
            geometry.append(mp.Prism(vertices = vertices2,
                                 height = L,
                                 material = Au))
            n_layer = int(resolution * t_structure)
            ### generate the 3D reflective layer structure by stacking blocks
            for k in range(n_layer):
                geometry.append(mp.Block(size = mp.Vector3(1/resolution, L, t_refl),
                                     center = mp.Vector3(x=-sx/2+dpml+dair+t_sub+t_co+1/resolution/2+k*1/resolution,
                                                         y=j*L,
                                                         z=sz/2-dpml-t_refl/2-k*1/resolution),
                                     material = Au))
                geometry.append(mp.Block(size = mp.Vector3(1/resolution, L, t_refl),
                                     center = mp.Vector3(x=-sx/2+dpml+dair+t_sub+t_co+1/resolution/2+k*1/resolution,
                                                         y=j*L,
                                                         z=-(sz/2-dpml-t_refl/2-k*1/resolution)),
                                     material = Au))
            geometry.append(mp.Block(size = mp.Vector3(t_electrode, L-t_structure, L-t_structure),
                                 center = mp.Vector3(sx/2-dpml-t_electrode/2, j*L),
                                 material = Au))  
    if plane == 'xy':
        cell_size = mp.Vector3(x=sx, y=sy)
    elif plane == 'yz':
        cell_size = mp.Vector3(y=sy, z=sz)
    pml_layer = [mp.PML(dpml)]                       ### set Perfectly matched layers as boundary layers
    beam_kdir = mp.Vector3(1, 0, 0)                  ### the propagating vector
    beam_w0 = L                                      ### the radius of the girdle
    beam_E0 = mp.Vector3(0, 1, 0)                    ### the polarization vector 
    beam_x0 = mp.Vector3(dair)                       ### distance from source to the substrate
    sources = [mp.GaussianBeamSource(mp.ContinuousSource(wavelength = wavelength, fwidth = fwidth),   ### using GaussianBeamSource
                                 size = mp.Vector3(0, L-2*t_refl, L-2*t_refl),
                                 center = mp.Vector3(-sx/2+dpml),
                                 beam_kdir = beam_kdir,
                                 beam_E0 = beam_E0,
                                 beam_x0 = beam_x0,
                                 beam_w0 = beam_w0)]
    symmetries = [mp.Mirror(mp.Y, phase = -1)]               ### Symmetries of the structure, set a proper symmetries will save your time
    sim = mp.Simulation(cell_size = cell_size,               
                    geometry = geometry,
                    resolution = resolution,
                    boundary_layers = pml_layer,
                    sources = sources,
                    symmetries = symmetries)
    plt.figure(dpi = 600)
    sim.plot2D()
    plt.xlabel('X/$\mu$m')
    plt.ylabel('Y/$\mu$m', rotation = 360)
    plt.show()
    sim.run(until = 200)
    QW_region = mp.Volume(size = mp.Vector3(periods*(t_well+t_br), L-2*t_refl, L-2*t_refl),             ### defining the active area of the target cell
                      center = mp.Vector3(sx/2-dpml-t_electrode-t_top-periods*(t_well+t_br)/2))
    Ex = sim.get_array(vol = QW_region, component = mp.Ex, cmplx = True)                                ### extract the Ex fields of the target cell
    eps = sim.get_array(vol = QW_region, component = mp.Dielectric)                                     ### extract the Dielectric constant matrix of the target cell
    (x, y, z, w) = sim.get_array_metadata(vol = QW_region)                                       ### extract the volume information
    energy_density = np.real(eps*(np.conj(Ex)*Ex))                                               ### calculate the energy density of the target cell
    energy = np.sum(w*energy_density)                                                            ### calculate the energy of the target cell
    Ix.append(energy)                                                                            ### save the energy
    print(energy)
    QW_region = mp.Volume(size = mp.Vector3(periods*(t_well+t_br), L-2*t_refl, L-2*t_refl),             ### defining the active arat of the adjacent cell
                      center = mp.Vector3(sx/2-dpml-t_electrode-t_top-periods*(t_well+t_br)/2, L))  
    Ex = sim.get_array(vol = QW_region, component = mp.Ex, cmplx = True)                                ### extract the Ex fields of the adjacent cell
    eps = sim.get_array(vol = QW_region, component = mp.Dielectric)                                     ### extract the Dielectric constant matrix of the adjacent cell
    (x, y, z, w) = sim.get_array_metadata(vol = QW_region)                                       ### extract the volume information
    energy_density = np.real(eps*(np.conj(Ex)*Ex))                                               ### calculate the energy density of the adjacent cell 
    energy = np.sum(w*energy_density)                                                            ### calculate the energy of the adjacent cell
    Ix_adj.append(energy)                                                                        ### save the energy
    Loop_num = Loop_num + 1                                                                      ### Loop monitor
    print("NO." + str(Loop_num) + "'s loop ")